"""
自定义数据接口辅助类
"""
from rest_framework.authentication import BaseAuthentication
from rest_framework.exceptions import AuthenticationFailed
from rest_framework.pagination import PageNumberPagination
from rest_framework.permissions import BasePermission

from common.models import UserToken


class CustomPagination(PageNumberPagination):
    """自定义分页"""

    page_size = 5
    page_size_query_param = 'size'
    max_page_size = 50


class CustomAuthentication(BaseAuthentication):
    """自定义身份认证"""

    def authenticate(self, request):
        token = request.META.get('HTTP_TOKEN', '')
        if token:
            user_token = UserToken.objects.filter(token=token).first()
            if user_token:
                return user_token.user, user_token
        raise AuthenticationFailed('请提供有效的用户身份标识')


class CustomPermission(BasePermission):
    """自定义权限类"""

    def has_permission(self, request, view):
        if getattr(request.user, 'roles', None):
            for role in request.user.roles.all().prefetch_related('privileges'):
                for priv in role.privileges.all():
                    if priv.method == request.method and request.path.startswith(priv.url):
                        return True
        return False
